{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true,
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "\n",
    "## Query Factor data and screen for different maximum and minimum thresholds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# import statements\n",
    "from typing import Dict, List, Union\n",
    "import pandas as pd\n",
    "from gs_quant.markets.factor import Factor\n",
    "import logging\n",
    "\n",
    "from gs_quant.session import GsSession, Environment\n",
    "\n",
    "from gs_quant.models.risk_model import FactorRiskModel, DataAssetsRequest, RiskModelUniverseIdentifierRequest as Identifier, \\\n",
    "    ReturnFormat\n",
    "\n",
    "# external users should substitute their client id and secret; please skip this step if using internal jupyterhub\n",
    "GsSession.use(Environment.PROD, client_id=None, client_secret=None, scopes=('read_product_data run_analytics',))\n",
    "\n",
    "\n",
    "# define a FactorThreshold class that takes in a minimum and maximum\n",
    "class FactorThreshold:\n",
    "    def __init__(self, factor: Factor, min_threshold: float = None, max_threshold: float = None):\n",
    "        self.factor = factor\n",
    "        self.min_threshold = min_threshold\n",
    "        self.max_threshold = max_threshold\n",
    "\n",
    "# define a helper function to do the heavy lifting\n",
    "def get_screened_universe(assets_zscore: pd.DataFrame,\n",
    "                          factors: List[FactorThreshold],\n",
    "                          use_absolute_value: bool = True,\n",
    "                          format: ReturnFormat = ReturnFormat.DATA_FRAME) -> Union[List[Dict], pd.DataFrame]:\n",
    "    \"\"\" Get universe zscores with exposures based on min and max threshold of exposures to factor list\n",
    "\n",
    "    :param assets_zscore: a pandas df of assets zscore to factors\n",
    "    :param factors: list of factors for screening\n",
    "    :param use_absolute_value: whether to threshold based on absolute value of factor exposure\n",
    "    :param format: which format to return the results in\n",
    "\n",
    "    :return: factor exposure for assets requested based on factor list input and max and min threshold values\n",
    "    \"\"\"\n",
    "\n",
    "    factors_in_zscore = assets_zscore.columns.tolist()\n",
    "    wrong_factors = [f.factor.name for f in factors if f.factor.name not in factors_in_zscore]\n",
    "    if wrong_factors:\n",
    "        logging.warning(f'Factors: {wrong_factors} not in the asset zscore input, will be filtered out..')\n",
    "\n",
    "    if assets_zscore.empty:\n",
    "        assets_zscore_return = {} if format == ReturnFormat.JSON else assets_zscore\n",
    "        return assets_zscore_return\n",
    "\n",
    "    factors_to_filter = [f for f in factors if f.factor.name not in wrong_factors]\n",
    "    for f in factors_to_filter:\n",
    "        if f.min_threshold is not None:\n",
    "            assets_zscore = assets_zscore[abs(assets_zscore[f.factor.name]) >= f.min_threshold] if use_absolute_value else \\\n",
    "                assets_zscore[assets_zscore[f.factor.name] >= f.min_threshold]\n",
    "        if f.max_threshold is not None:\n",
    "            assets_zscore = assets_zscore[abs(assets_zscore[f.factor.name]) <= f.max_threshold] if use_absolute_value else \\\n",
    "                assets_zscore[assets_zscore[f.factor.name] <= f.max_threshold]\n",
    "    if format == ReturnFormat.JSON:\n",
    "        assets_zscore = assets_zscore.to_dict()\n",
    "    return assets_zscore\n",
    "\n",
    "\n",
    "# get the model and dates available\n",
    "model_id = 'BARRA_USMEDS'\n",
    "factor_model = FactorRiskModel.get(model_id)\n",
    "\n",
    "# get multiple measures across a date range for a universe specified\n",
    "dates = factor_model.get_dates()\n",
    "latest_date = dates.pop()\n",
    "\n",
    "total_factors = factor_model.get_many_factors(start_date=latest_date, end_date=latest_date)\n",
    "style_factors = [f for f in total_factors if f.type != 'Category' and f.category == \"Style\"]\n",
    "print([f.name for f in style_factors])\n",
    "\n",
    "# set maximum and minimum zscore threshold for the factors\n",
    "growth = factor_model.get_factor('Growth')\n",
    "beta = factor_model.get_factor('Beta')\n",
    "\n",
    "style_factor_thresholds = [\n",
    "    FactorThreshold(factor=growth, min_threshold=0, max_threshold=1),\n",
    "    FactorThreshold(factor=beta, min_threshold=-1, max_threshold=0)\n",
    "]\n",
    "\n",
    "\n",
    "assets = DataAssetsRequest(Identifier.sedol, [])\n",
    "assets_to_zscore = factor_model.get_universe_factor_exposure(\n",
    "    start_date=latest_date,\n",
    "    end_date=latest_date,\n",
    "    assets=assets,\n",
    "    get_factors_by_name=True\n",
    ").reset_index(level=1, drop=True)\n",
    "\n",
    "print(assets_to_zscore.head())\n",
    "\n",
    "\n",
    "screened = get_screened_universe(assets_to_zscore, style_factor_thresholds)\n",
    "print(screened.head())"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}